import pandas as pd
from py2neo import Node, Relationship, Graph, NodeMatcher, RelationshipMatcher
import pymysql
from pymysql.cursors import DictCursor
from tqdm import tqdm


def connect_db():
    db = pymysql.connect (user="root", password="sufeFin", host="127.0.0.1", database="sufeFin", port=8306)
    return db


# 创建节点
def CreateNode(m_graph, m_label, m_attrs):
    try:
        m_n = f"""_.name="{m_attrs['name']}" """
        matcher = NodeMatcher (m_graph)
        re_value = matcher.match (m_label).where (m_n).first ()
        # print(re_value)
        if re_value is None:
            m_mode = Node (m_label, **m_attrs)
            n = graph.create (m_mode)
            return n
    except Exception as e:
        print(e)
    return None


# 查询节点
def MatchNode(m_graph, m_label, m_attrs):
    try:
        m_n = f"""_.name="{m_attrs['name']}" """
        matcher = NodeMatcher (m_graph)
        re_value = matcher.match (m_label).where (m_n).first ()
        return re_value
    except Exception as e:
        print(e)
        return None


# 创建关系
def CreateRelationship(m_graph, m_label1, m_attrs1, m_label2, m_attrs2, m_r_name):
    try:
        reValue1 = MatchNode (m_graph, m_label1, m_attrs1)
        reValue2 = MatchNode (m_graph, m_label2, m_attrs2)
        if reValue1 is None or reValue2 is None:
            return False
        m_r = Relationship (reValue1, m_r_name, reValue2)
        n = graph.create (m_r)
        return n
    except Exception as e:
        print(e)
        return None


if __name__ == '__main__':
    db = connect_db ()
    with db.cursor () as cursor:
        cursor.execute (""" select distinct name from holder """)
        res = cursor.fetchall ()
        companies = [o[0] for o in res]
        cursor.execute (""" select distinct cpafirm.cpafirm from cpafirm """)
        res = cursor.fetchall ()
        cpafirm = [o[0] for o in res]
        cursor.execute (""" select distinct lawfirm.lawfirm from lawfirm """)
        res = cursor.fetchall ()
        lawfirm = [o[0] for o in res]
        cursor.execute (
            """ select distinct company_name,shareholder_name,share_ratio  from top10_holder where not( ISNULL(company_name) or  ISNULL(shareholder_name) or ISNULL(share_ratio) )""")
        res = cursor.fetchall ()
        top_holder_data = res
        cursor.execute (""" select distinct secCode,secName  from secCodes """)
        res = cursor.fetchall ()
        code_name_dic = {c: n for c, n in res}
        cursor.execute (
            """ select distinct full_name,cpafirm  from company_info""")
        res = cursor.fetchall ()
        cpa_record = res
        cursor.execute (
            """ select distinct full_name,lawfirm  from company_info""")
        res = cursor.fetchall ()
        law_record = res
    # data = pd.DataFrame({"名称":Names,"字段":things,"方式":action})
    invest_data = pd.read_csv ("financial_data/processed_long_term_invest.csv")
    invest_data['company'] = invest_data['证券代码'].map (lambda x: code_name_dic.get (x))
    invest_data = invest_data[['company', '长期投资项目', '权益比率']]
    invest_data.fillna ('未知')
    invest_data.columns = ['company', 'invest_company', 'invest_ratio']

    username = 'neo4j'
    password = 'sufeFin'
    graph = Graph ('http://localhost:7474', username=username, password=password)

    label1 = "holder"
    label2 = "cpa"
    label3 = "law"

    # for i, j, r in tqdm (top_holder_data, desc="holded"):
    #     if i is None or j is None or r is None:
    #         continue
    #     if i.strip () == "" or j.strip () == "":
    #         continue
    #     # company
    #     attr1 = {"name": i}
    #     CreateNode (graph, label1, attr1)
    #     # shareholder
    #     attr2 = {"name": j}
    #     CreateNode (graph, label1, attr2)
    #
    #     m_r_name = "被持有"
    #     reValue = CreateRelationship (graph, label1, attr1, label1, attr2, m_r_name)

    for index,(i, j, r) in tqdm (invest_data.iterrows (), desc="invest"):
        if i is None or j is None or r is None:
            continue
        if i.strip () == "" or j.strip () == "":
            continue
        # company
        attr1 = {"name": i}
        CreateNode (graph, label1, attr1)
        # invested company
        attr2 = {"name": j}
        CreateNode (graph, label1, attr2)

        m_r_name = "长期投资"
        reValue = CreateRelationship (graph, label1, attr1, label1, attr2, m_r_name)

    for i, j in tqdm (cpa_record, desc="cpa"):
        if i is None or j is None:
            continue
        if i.strip () == "" or j.strip () == "":
            continue
        # company
        attr1 = {"name": i}
        CreateNode (graph, label1, attr1)
        # invested company
        attr2 = {"name": j}
        CreateNode (graph, label2, attr2)

        m_r_name = "被审计"
        reValue = CreateRelationship (graph, label1, attr1, label2, attr2, m_r_name)

    for i, j in tqdm (law_record, desc="law"):
        if i is None or j is None:
            continue
        if i.strip () == "" or j.strip () == "":
            continue
        # company
        attr1 = {"name": i}
        CreateNode (graph, label1, attr1)
        # invested company
        attr2 = {"name": j}
        CreateNode (graph, label3, attr2)

        m_r_name = "法律咨询"
        reValue = CreateRelationship (graph, label1, attr1, label3, attr2, m_r_name)
